package algorithm.ordered_table;

//非积压结构
//查询、添加、删除节点都是O（log N）
public class AVLTreeMap {
    public static class AVLNode<K extends Comparable<? super K>, V> {
        private K key;
        private V value;
        private AVLNode<K, V> l;
        private AVLNode<K, V> r;
        private int h;

        public AVLNode(K key, V value) {
            this.key = key;
            this.value = value;
            h = 1;
        }
    }

    public static class MyAVLTreeMap<K extends Comparable<? super K>, V> {
        private AVLNode<K, V> root;
        private int size;

        public MyAVLTreeMap() {
            root = null;
            size = 0;
        }

        private AVLNode<K, V> rightRotate(AVLNode<K, V> cur) {
            AVLNode<K, V> left = cur.l;
            cur.l = left.r;
            left.r = cur;
            cur.h = Math.max(cur.l == null ? 0 : cur.l.h, cur.r == null ? 0 : cur.r.h) + 1;
            left.h = Math.max(left.l == null ? 0 : left.l.h, left.r == null ? 0 : left.r.h) + 1;
            return left;
        }

        private AVLNode<K, V> leftRotate(AVLNode<K, V> cur) {
            AVLNode<K, V> right = cur.r;
            cur.r = right.l;
            right.l = cur;
            cur.h = Math.max(cur.l == null ? 0 : cur.l.h, cur.r == null ? 0 : cur.r.h) + 1;
            right.h = Math.max(right.l == null ? 0 : right.l.h, right.r == null ? 0 : right.r.h) + 1;
            return right;
        }

        //AVL树保持平衡的策略只要进行单次或者两次的左旋或者右旋即可
        private AVLNode<K, V> maintain(AVLNode<K, V> cur) {
            if (cur == null)
                return null;
            int leftHeight = cur.l == null ? 0 : cur.l.h;
            int rightHeight = cur.r == null ? 0 : cur.r.h;
            if (Math.abs(leftHeight - rightHeight) > 1) {
                if (leftHeight > rightHeight) {
                    int leftLeftHeight = cur.l != null ? cur.l.l == null ? 0 : cur.l.l.h : 0;
                    int leftRightHeight = cur.l != null ? cur.l.r == null ? 0 : cur.l.r.h : 0;
                    if (leftLeftHeight >= leftRightHeight)
                        cur = rightRotate(cur);
                    else {
                        cur.l = leftRotate(cur.l);
                        cur = rightRotate(cur);
                    }
                } else {
                    int rightLeftHeight = cur.r != null ? cur.r.l == null ? 0 : cur.r.l.h : 0;
                    int rightRightHeight = cur.r != null ? cur.r.r == null ? 0 : cur.r.r.h : 0;
                    if (rightRightHeight >= rightLeftHeight)
                        cur = leftRotate(cur);
                    else {
                        cur.r = rightRotate(cur.r);
                        cur = leftRotate(cur);
                    }
                }
            }
            return cur;
        }

        private AVLNode<K, V> add(AVLNode<K, V> cur, K key, V value) {
            if (cur == null)
                return new AVLNode<>(key, value);
            if (cur.key.compareTo(key) > 0) {
                cur.l = add(cur.l, key, value);
            } else {
                cur.r = add(cur.r, key, value);
            }
            cur.h = Math.max(cur.l == null ? 0 : cur.l.h, cur.r == null ? 0 : cur.r.h) + 1;
            return maintain(cur);
        }

        private AVLNode<K, V> delete(AVLNode<K, V> cur, K key) {
            if (cur.key.compareTo(key) > 0)
                cur.l = delete(cur.l, key);
            else if (cur.key.compareTo(key) < 0)
                cur.r = delete(cur.r, key);
            else {
                if (cur.l == null && cur.r == null)
                    cur = null;
                else if (cur.l == null && cur.r != null)
                    cur = cur.r;
                else if (cur.l != null && cur.r == null)
                    cur = cur.l;
                else {
                    AVLNode<K, V> des = cur.l;
                    while (des.r != null)
                        des = des.r;
                    cur.l = delete(cur.l, des.key);
                    cur.key = des.key;
                    cur.value = des.value;
                }
            }
            if (cur != null)
                cur.h = Math.max(cur.l == null ? 0 : cur.l.h, cur.r == null ? 0 : cur.r.h) + 1;
            return maintain(cur);
        }

        private AVLNode<K, V> findLastIndex(K key) {
            AVLNode<K, V> cur = root;
            AVLNode<K, V> pre = root;
            while (cur != null) {
                if (key.compareTo(cur.key) == 0)
                    return cur;
                pre = cur;
                cur = cur.key.compareTo(key) > 0 ? cur.l : cur.r;
            }
            return pre;
        }

        private AVLNode<K, V> findLastNoSmallIndex(K key) {//找到第一个比他大或者等于他的节点
            AVLNode<K, V> cur = root;
            AVLNode<K, V> ans = null;
            while (cur != null) {
                if (key.compareTo(cur.key) == 0) {
                    ans = cur;
                    break;
                } else if (key.compareTo(cur.key) < 0) {//从当前节点向左走说明当前节点以及当前节点的左树都比当前查找的key值小
                    ans = cur;
                    cur = cur.l;
                } else {
                    cur = cur.r;
                }
            }
            return ans;
        }

        private AVLNode<K, V> findLastNoBigIndex(K key) {//找到最后一个比他小或者等于他的节点
            AVLNode<K, V> cur = root;
            AVLNode<K, V> ans = null;
            while(cur != null){
                if(key.compareTo(cur.key) == 0){
                    ans = cur;
                    break;
                }else if(key.compareTo(cur.key) > 0){//从当前节点向右走说明当前节点以及当前节点的左树都比当前查找的key值小
                    ans = cur;
                    cur = cur.r;
                }else
                    cur = cur.l;
            }
            return ans;
        }

        public boolean containsKey(K key) {
            if(key == null)
                return false;
            AVLNode<K, V> lastIndex = findLastIndex(key);
            if(lastIndex != null && lastIndex.key.compareTo(key) == 0){
                return true;
            }else {
                return false;
            }
        }

        public void put(K key, V value) {
            if(key != null){
                AVLNode<K, V> lastIndex = findLastIndex(key);
                if(lastIndex != null && lastIndex.key.compareTo(key) == 0)
                    lastIndex.value = value;
                else {
                    AVLNode<K, V> root = add(this.root, key, value);
                    this.root = root;
                    this.size++;
                }
            }
        }

        public void remove(K key) {
            if(key != null){
                if(containsKey(key)){
                    AVLNode<K, V> node = delete(root, key);
                    this.root = node;
                    this.size--;
                }
            }
        }

        public V get(K key) {
            if(key != null){
                AVLNode<K, V> lastIndex = findLastIndex(key);
                return lastIndex != null && lastIndex.key.compareTo(key) == 0 ? lastIndex.value : null;
            }
            return null;
        }
    }

    public static void main(String[] args) {
        MyAVLTreeMap<Integer, String> map = new MyAVLTreeMap<>();
        map.put(4, "xff");
        map.put(3, "gg");
        map.put(2, "ll");
        map.put(1, "uy");
        String s = map.get(4);
//        System.out.println(map.containsKey(0));
        System.out.println(s);
        map.remove(4);
        System.out.println(map.get(3));
    }

}
